from sklearn.preprocessing import OneHotEncoder
import numpy as np

'''
about converting DataFrame to training format 
'''

def filter_samples(df, x_fields, y_field, city = "101"):
    x_fields_condition = ["df['{}'].isnull() == False".format(x) for x in x_fields]
    df_train = df[(df[y_field] < 999990) &    # reject missing y
                  (df[y_field].isnull().values == False) &  # reject NaN y
                  eval('({})'.format(')&('.join(x_fields_condition) ))   #reject NaN x
                  ]
    x_part, y_part = df_train[x_fields], df_train[y_field]
    return x_part, y_part


def train_to_array(df_categories):
    ohe1 = OneHotEncoder(handle_unknown='ignore')
    ar = ohe1.fit_transform(df_categories).toarray()
    return ohe1, np.array(ar)


def test_to_array(df_categories, ohe):
    ar = ohe.transform(df_categories).toarray()
    return np.array(ar)

def merge_matrix(aa, bb):
    new_matrix = []
    assert len(aa) == len(bb)
    for i in range(len(aa)):
        new_matrix.append(np.hstack((aa[i], bb[i])))
    return np.array(new_matrix)


def convert_traindf_to_matrix(x_train_df, category_fields, numeric_fields):
    df_category = x_train_df[category_fields]
    df_numeric = x_train_df[numeric_fields]
    onehot_encoder, x_ar = train_to_array(df_category)
    x_ar2 = np.array(df_numeric)
    print(x_ar.shape, x_ar2.shape , end =" > ")
    x_train = merge_matrix(x_ar, x_ar2)
    return x_train, onehot_encoder

def convert_testdf_to_matrix(x_test_df, category_fields, numeric_fields, onehot_encoder):
    df_category = x_test_df[category_fields]
    df_numeric = x_test_df[numeric_fields]
    x_ar = test_to_array(df_category, onehot_encoder)
    x_ar2 = np.array(df_numeric)
    print(x_ar.shape, x_ar2.shape, end=" > ")
    x_test = merge_matrix(x_ar, x_ar2)
    return x_test


def training(x_train, y_train, model):
    model.fit(x_train, y_train)
    return model
    # score = model.score(x_train, y_train)
    # print(" score on training set:{}".format(score))


def predicting(x_test, model):
    return model.predict(x_test)

# def run_train(x_train_df, y_train, category_fields, numeric_fields, model):
#     df_category = x_train_df[category_fields]
#     df_numeric = x_train_df[numeric_fields]
#     onehot_encoder, x_ar = train_to_array(df_category)
#     x_ar2 = np.array(df_numeric)
#     print(x_ar.shape, x_ar2.shape , end =" > ")
#     x_train = merge_matrix(x_ar, x_ar2)
#     print(x_train.shape, y_train.shape)
#     model.fit(x_train, y_train)
#     return model, onehot_encoder